#include "call_rates.hpp"

void do_callrates(int argc, char **argv)
{
    prog_options options = parse_callrates_options(argc, argv);

    mapping_type t = get_mapping_type(options.map_file.c_str());

    if(t & COLOR)
    {
        // TODO: UGLY. 
        options.old_bin_cnt = 36;
        do_color_callrates(options);
    }
    else 
    {
        // TODO: UGLY. 
        options.old_bin_cnt = 96;
        do_seq_callrates(options);
    }
}

void do_color_callrates(prog_options &options)
{
    array_type5 calls;
    cand_blocks_t cand_blocks;
    cand_vector_t snp_cands;
    
    if(options.has_cand)
    {
        read_snp_cand_file(options.cand_file.c_str(), cand_blocks, snp_cands);
        find_call_rates(options, calls, cand_blocks, snp_cands);
    }
    else
    {
        find_call_rates_no_cand(options, calls);
    }

    write_call_counts(std::cout, calls);
}

void do_seq_callrates(prog_options &options)
{
    array_type4 calls;
    cand_blocks_t cand_blocks;
    cand_vector_t snp_cands;

    
    if(options.has_cand)
    {
        read_snp_cand_file(options.cand_file.c_str(), cand_blocks, snp_cands);
        find_seq_call_rates(options, calls, cand_blocks, snp_cands);
    }
    else
    {
        find_seq_call_rates_no_cand(options, calls);
    }

    write_seq_call_counts(std::cout, calls);
}

void process_color_counts_to_rates(array_type5 &call_counts, array_type5 &call_rates, int obc, int nbc, int read_length, boost::multi_array<int, 2> &qual_bins)
{
    // rebin quals.
    rebin_quals(call_counts, read_length, obc, nbc, qual_bins);
    
    // use the new bin boundaries to collapse the call_rates matrix.
    rebin_calls(call_counts, qual_bins, read_length, obc, nbc, call_rates);

    // convert the counts to probabilities.
    convert_call_rates_to_probs(call_rates, read_length, nbc);

}

void process_seq_counts_to_rates(array_type4 &call_counts, array_type4 &call_rates, int obc, int nbc, int read_length, boost::multi_array<int, 2> &qual_bins)
{
    // rebin quals.
    rebin_seq_quals(call_counts, read_length, obc, nbc, qual_bins);
    
    // use the new bin boundaries to collapse the call_rates matrix.
    rebin_seq_calls(call_counts, qual_bins, read_length, obc, nbc, call_rates);

    // convert the counts to probabilities.
    convert_seq_call_rates_to_probs(call_rates, read_length, nbc);

}

void fill_call_priors(prog_options &options, array_type5 &calls)
{
    array_type5::extent_gen extents;

    // TODO: quality filter?
    unsigned int read_length = get_read_length(options.map_file.c_str());

    calls.resize( extents[read_length][options.old_bin_cnt][4][4][4] );
    
    unsigned int i,j,k,l,m;
    for(i = 0; i < read_length; i++)
        for(j = 0; j < options.old_bin_cnt; j++)
        {
            for(k = 0; k < 4; k++)
            {
                for(l = 0; l < 4; l++)
                    for(m = 0; m < 4; m++)
                        calls[i][j][k][l][m] = 1.0;
                
                calls[i][j][k][k][0] = 10;
            }
            
            calls[i][j][0][3][3] = 10;
            calls[i][j][3][0][3] = 10;
            calls[i][j][1][2][3] = 10;
            calls[i][j][2][1][3] = 10;

            calls[i][j][0][1][1] = 10;
            calls[i][j][1][0][1] = 10;
            calls[i][j][2][3][1] = 10;
            calls[i][j][3][2][1] = 10;

            calls[i][j][0][2][2] = 100;
            calls[i][j][2][0][2] = 100;
            calls[i][j][1][3][2] = 100;
            calls[i][j][3][1][2] = 100;
        }
}


void find_call_rates_no_cand(prog_options &options, array_type5 &calls)
{
    // priors for call rate matrix.  All transitions are scored with 100 if
    // they are expected, and 1 if not.  These pseudo counts help to provide
    // good numbers even in areas where there is not a lot of read coverage.
    //

    fill_call_priors(options, calls);

    unsigned int total = 0;
    int skipped_uniq = 0;
    std::string dna_ref;


    ColorMappingParser _parser(options.map_file.c_str());

    ColorRead r;
    while( _parser.next_read(r) && total < options.read_count )
    {

        total++;

        if( options.unique && r.map_locs.size() > 1 )
        {
            skipped_uniq++;
            continue;
        }

        for(unsigned int i = 0; i < r.map_locs.size(); i++)
        {

            MapLoci m = r.map_locs[i];

            // TODO: screen for quality
            dna_ref = color_to_seq(m.seq);

            // The transition between genome position i and genome position i +
            // 1 gave the color at read position i+1
            //
            // The call rate matrix holds the counts of observed transitions.
            // 
            // calls[ p ][ q ][ b1 ][ b2 ][ c ]
            // 
            // is the number of occurances at position p in the read, with
            // quality q at that position, where the genome sequence at p-1 is
            // b1 and p is b2 and the color c was produced.
            //
            // The strand does not affect this computation.
            //
            // We walk over the genome and increment the counts.
            for(unsigned int j = 0; j < dna_ref.size()-1; j++)
                calls[j+1][r.qual[j+1]][cswap[dna_ref[j]]][cswap[dna_ref[j+1]]][iswap[r.seq[j+1]]]++;
        } 
    }
}

void find_call_rates(prog_options &options, array_type5 &calls, cand_blocks_t &cand_blocks, cand_vector_t &snp_cands)
{

    fill_call_priors(options, calls);

    unsigned int total = 0;
    int skipped_uniq = 0;
    std::string dna_ref;


    ColorMappingParser _parser(options.map_file.c_str());

    ColorRead r;
    while( _parser.next_read(r) && total < options.read_count )
    {

        total++;

        if( options.unique && r.map_locs.size() > 1 )
        {
            skipped_uniq++;
            continue;
        }

        for(unsigned int i = 0; i < r.map_locs.size(); i++)
        {

            MapLoci m = r.map_locs[i];
            
            if(cand_blocks.find(m.id) == cand_blocks.end())
            {
                // ignore-missing is always on.
                continue;
            }

            // TODO: screen for quality
            dna_ref = color_to_seq(m.seq);

            // The transition between genome position i and genome position i +
            // 1 gave the color at read position i+1
            //
            // The call rate matrix holds the counts of observed transitions.
            // 
            // calls[ p ][ q ][ b1 ][ b2 ][ c ]
            // 
            // is the number of occurances at position p in the read, with
            // quality q at that position, where the genome sequence at p-1 is
            // b1 and p is b2 and the color c was produced.
            //
            // The strand does not affect this computation.
            //
            // We walk over the genome and increment the counts.
            for(unsigned int j = 0; j < dna_ref.size()-1; j++)
                if( ! cand_blocks[m.id][m.pos+j] )
                    calls[j+1][r.qual[j+1]][cswap[dna_ref[j]]][cswap[dna_ref[j+1]]][iswap[r.seq[j+1]]]++;
        } }
}

void fill_call_priors_seq(prog_options &options, array_type4 &calls)
{

    array_type4::extent_gen extents;

    unsigned int read_length = get_read_length(options.map_file.c_str());

    calls.resize( extents[read_length][options.old_bin_cnt][4][4] );
    
    // priors for call rate matrix.  All transitions are scored with 100 if
    // they are expected, and 1 if not.  These pseudo counts help to provide
    // good numbers even in areas where there is not a lot of read coverage.
    
    unsigned int i,j,k,l;
    for(i = 0; i < read_length; i++)
        for(j = 0; j < options.old_bin_cnt; j++)
        {
            for(k = 0; k < 4; k++)
            {
                for(l = 0; l < 4; l++)
                        calls[i][j][k][l] = 1.0;
                
                calls[i][j][k][k] = 100;
            }
        }
}

void find_seq_call_rates_no_cand(prog_options &options, array_type4 &calls)
{
    fill_call_priors_seq(options, calls);

    unsigned int total = 0;
    int skipped_uniq = 0;
    std::string dna_ref;

    mapping_type t = get_mapping_type(options.map_file.c_str());

    if(t & PERM_FILE)
    {
        SequenceMappingParser _parser(options.map_file.c_str());

        SeqRead r;
        while( _parser.next_read(r) && total < options.read_count )
        {

            total++;

            if( options.unique && r.map_locs.size() > 1 )
            {
                skipped_uniq++;
                continue;
            }

            for(unsigned i = 0; i < r.map_locs.size(); i++)
            {

                MapLoci m = r.map_locs[i];

                // The call rate matrix holds the counts of observed transitions.
                // 
                // calls[ p ][ q ][ b1 ][ b2 ]
                // 
                // is the number of occurances at position p in the read, with
                // quality q at that position, where the genome sequence is b1 and the read sequence is b2.
                //
                // The strand does not affect this computation.
                //
                // We walk over the genome and increment the counts.
                for(unsigned int j = 0; j < r.seq.size(); j++)
                    calls[j][r.qual[j]][cswap[m.seq[j]]][cswap[r.seq[j]]]++;
            } 
        }
    }
    else if(t & SAM_FILE)
    {
        std::map<std::string, std::string> seqs;
        parse_fasta(options.fasta_file.c_str(), seqs);
        SamSequenceMappingParser _parser(options.map_file.c_str(), seqs, t);

        SeqRead r;
        while( _parser.next_read(r) && total < options.read_count )
        {

            total++;

            if( options.unique && r.map_locs.size() > 1 )
            {
                skipped_uniq++;
                continue;
            }

            for(unsigned i = 0; i < r.map_locs.size(); i++)
            {

                MapLoci m = r.map_locs[i];

                // The call rate matrix holds the counts of observed transitions.
                // 
                // calls[ p ][ q ][ b1 ][ b2 ]
                // 
                // is the number of occurances at position p in the read, with
                // quality q at that position, where the genome sequence is b1 and the read sequence is b2.
                //
                // The strand does not affect this computation.
                //
                // We walk over the genome and increment the counts.
                for(unsigned int j = 0; j < r.seq.size(); j++)
                    calls[j][r.qual[j]][cswap[m.seq[j]]][cswap[r.seq[j]]]++;
            } 
        }
    }
    else
        error("Bad File type");
}

void find_seq_call_rates(prog_options &options, array_type4 &calls, cand_blocks_t &cand_blocks, cand_vector_t &snp_cands)
{

    fill_call_priors_seq(options, calls);

    unsigned int total = 0;
    int skipped_uniq = 0;

    mapping_type t = get_mapping_type(options.map_file.c_str());

    if(t & PERM_FILE)
    {
        SequenceMappingParser _parser(options.map_file.c_str());

        SeqRead r;
        while( _parser.next_read(r) && total < options.read_count )
        {

            total++;

            if( options.unique && r.map_locs.size() > 1 )
            {
                skipped_uniq++;
                continue;
            }

            for(unsigned int i = 0; i < r.map_locs.size(); i++)
            {

                MapLoci m = r.map_locs[i];
                
                if(cand_blocks.find(m.id) == cand_blocks.end())
                {
                    // ignore-missing is always on.
                    continue;
                }

                // The call rate matrix holds the counts of observed transitions.
                // 
                // calls[ p ][ q ][ b1 ][ b2 ]
                // 
                // is the number of occurances at position p in the read, with
                // quality q at that position, where the genome sequence is b1 and the read sequence is b2.
                //
                // The strand does not affect this computation.
                //
                // We walk over the genome and increment the counts.
                for(unsigned int j = 0; j < r.seq.size(); j++)
                    calls[j][r.qual[j]][cswap[m.seq[j]]][cswap[r.seq[j]]]++;
            } 
        }
    }
    else if(t & SAM_FILE)
    {
        std::map<std::string, std::string> seqs;
        parse_fasta(options.fasta_file.c_str(), seqs);
        SamSequenceMappingParser _parser(options.map_file.c_str(), seqs, t);

        SeqRead r;
        while( _parser.next_read(r) && total < options.read_count )
        {

            total++;

            if( options.unique && r.map_locs.size() > 1 )
            {
                skipped_uniq++;
                continue;
            }

            for(unsigned int i = 0; i < r.map_locs.size(); i++)
            {

                MapLoci m = r.map_locs[i];
                
                if(cand_blocks.find(m.id) == cand_blocks.end())
                {
                    // ignore-missing is always on.
                    continue;
                }

                // The call rate matrix holds the counts of observed transitions.
                // 
                // calls[ p ][ q ][ b1 ][ b2 ]
                // 
                // is the number of occurances at position p in the read, with
                // quality q at that position, where the genome sequence is b1 and the read sequence is b2.
                //
                // The strand does not affect this computation.
                //
                // We walk over the genome and increment the counts.
                for(unsigned int j = 0; j < r.seq.size(); j++)
                    calls[j][r.qual[j]][cswap[m.seq[j]]][cswap[r.seq[j]]]++;
            } 
        }
    }
    else
        error("Bad file type");
}
